import json, os
from .cocostuff_loader import CocoSceneGraphDataset
from .vg import VgSceneGraphDataset


def get_dataset(dataset='coco', img_size=0, left_right_flip=False, train=False, dir_dataset = './datasets', max_objects_per_image=8):
    dir_dataset = os.path.abspath(dir_dataset)
    if train:
        if dataset == "coco":
            return_dataset = CocoSceneGraphDataset(image_dir=f'{dir_dataset}/coco/train2017/',
                                            instances_json=f'{dir_dataset}/coco/annotations/instances_train2017.json',
                                            stuff_json=f'{dir_dataset}/coco/annotations/stuff_train2017.json',
                                            stuff_only=True, image_size=(img_size, img_size), left_right_flip=left_right_flip, max_objects_per_image=max_objects_per_image)
        elif dataset == 'vg':
            with open(f"{dir_dataset}/vg/vocab.json", "r") as read_file:
                vocab = json.load(read_file)
            return_dataset = VgSceneGraphDataset(vocab=vocab, h5_path=f'{dir_dataset}/vg/train.h5',
                                        image_dir=f'{dir_dataset}/vg/images/',
                                        image_size=(img_size, img_size), max_objects=10, left_right_flip=left_right_flip)
    else:
        if dataset == 'coco':
            return_dataset = CocoSceneGraphDataset(image_dir=f'{dir_dataset}/coco/val2017/',
                                            instances_json=f'{dir_dataset}/coco/annotations/instances_val2017.json',
                                            stuff_json=f'{dir_dataset}/coco/annotations/stuff_val2017.json',
                                            stuff_only=True, image_size=(img_size, img_size), left_right_flip=left_right_flip, max_objects_per_image=max_objects_per_image)
        elif dataset == 'vg':
            with open(f"{dir_dataset}/vg/vocab.json", "r") as read_file:
                vocab = json.load(read_file)
            return_dataset = VgSceneGraphDataset(vocab=vocab,
                                        h5_path=f'{dir_dataset}/vg/val.h5',
                                        image_dir=f'{dir_dataset}/vg/images/',
                                        image_size=(img_size, img_size), 
                                        left_right_flip=left_right_flip, max_objects=10)
    return return_dataset